fix(infer): normalize parameter shorthand before batching - #5857

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/pt/infer/deep_eval.py`:
- Around line 548-555: Update eval_embedding to call _standardize_fparam_aparam
on fparam and aparam before dispatching to self._eval_func, using the embedding
evaluation’s frame count, natoms, and model parameter dimensions. Pass the
standardized values to the backend so shorthand inputs are batched by frames
correctly.

The core change is right, and putting the ladder in one shared helper called before AutoBatchSize is the correct design. I checked it rather than assuming: full-shape inputs are unchanged across a wide sweep of dim_fparam / dim_aparam / nframes / auto_batch_size combinations, the public path stays idempotent (the flattened 2-D form re-enters on the nframes*natoms*dim rung and only reshapes, no double-tiling), the three size-dispatch collision families (nframes==1, natoms==1, both) produce identical arrays in every colliding branch, and the spin route gets the real-atom natoms that _eval_model_spin expects. It also fixes the long-standing "got wrong size of frame param" typo in the aparam branch.

I ran the new test_embedding.py case against the pre-fix tree and it fails with ValueError: cannot reshape array of size 1 into shape (1,7,1), so that one is a genuine regression test - two frames, a forced per-frame split, and np.linspace(0.1, 0.7, natoms) so aparam actually varies per atom. Good test.

My comments are all about coverage and scope, not the helper. Six inline. The two I would weigh most are the ones about what the tests can actually detect: test_io.py uses uniform values, and on a GPU runner neither new test performs the split it is built around.

Note the head moved while I was reviewing (bbfc000ac2b1 -> 8d80b8579); everything below is against 8d80b8579.

Smaller notes, not worth their own threads:

1. The TF v1 backend has the same defect and is worse than the others: deepmd/tf/infer/deep_eval.py passes fparam=/aparam= as kwargs into execute_all (which slices kwargs too), and its shorthand ladder lives inside _prepare_feed_dict, i.e. per batch, after slicing. With one frame per batch a (natoms, dim_aparam) array is sliced on its atom axis to (1, dim_aparam), which then matches the size == dim_aparam rung and tiles atom 0's value onto every atom - wrong numbers, no exception. Out of this PR's stated scope, but "tensorflow" is in the same test loop and its exclusion is the undocumented one.

2. The PR body says it covers "both normal and spin backend routes". The placement before the spin branch is correct, but no test in the diff constructs a spin model or passes spin=, so _eval_model_spin is never exercised with shorthand.

3. test_embedding.py's assertEqual(backend.auto_batch_size.current_batch_size, natoms) fails outright if DP_INFER_BATCH_SIZE is exported, since AutoBatchSize.__init__ lets that variable override the constructor argument - and if it is set, the test silently stops forcing the split while still passing.

4. The helper unit tests never pass fparam and aparam together, which is the shape every production call site uses. The two branches are independent so nothing is likely hiding there.

5. _standard_input builds the canonical 3-D aparam and immediately re-flattens it to 2-D. Every consumer reshapes from .size regardless of incoming rank, so nothing needs the 2-D form any more; dropping the re-flatten would make the internal boundary unambiguous.

6. The aparam error text changes from "frame param" to "atomic param" (a real fix), but deepmd/tf/infer/deep_eval.py still carries the old wording in its own duplicate check, so the two entry points now disagree. Worth a line in the PR description since it is user-visible.

Thanks -- I went through all six of my earlier threads against HEAD f785c9e0b rather than the replies, and they are genuinely addressed:

• The consistency test no longer compares byte-identical frames. Distinct coordinates (+0.125), a perturbed box, linspace parameters and distinct per-frame rows mean a mis-tiling now changes a number instead of coinciding.
• The batching half of the rationale is restored: each case builds a fresh DeepEval(model_file, auto_batch_size=natoms) and issues the shorthand call first, with a separate unbatched evaluator for the reference, so GPU batch-size growth cannot silently remove the split. Relaxing to 1e-10 for the cross-batch comparison is the right call given differing GEMM shapes.
• The allowlist gap is closed the better way -- dpmodel and tensorflow are exercised rather than excluded, and Paddle's absence now carries a stated reason instead of being invisible.
• The pd hunk has executing coverage again, and the PR body is corrected to say plainly that it runs in the Paddle environment rather than locally.
• _standard_input itself now delegates to the helper and re-flattens to the historical 2-D ABI, so the duplication I was worried about is genuinely gone: one definition serving the public wrapper and every adapter.
• The helper's numpydoc block covers the parameters and both return shapes, including the 3-D aparam and the flattening note.

I also checked the two things that would worry me about a refactor of this shape. An exhaustive sweep of the new helper against master's _standard_input ladder over nframes x natoms x dim_fparam x dim_aparam x input size x candidate shapes produced no difference in value, shape, or raise behaviour -- the .size dispatch is ambiguous only in the cases where both branches yield the same array, exactly as before. And on the regression side I ran your new pt_expt adapter test against unfixed code: it fails there, so it does prove the fix.

One gap left, inline, on the adapter I asked about last time. I would rather see it closed here than tracked separately, since it is the same contract and the same file.

Two smaller notes, neither needing action in this PR. deepmd/tf/infer/deep_eval.py#L850-L862 still holds the old private ladder; both of its callers now receive normalized input so it is unreachable in practice, but it still carries the copy-paste bug your helper fixes -- the aparam branch raises "got wrong size of frame param". Worth deleting or reducing to an assertion while the context is fresh. And PretrainedDeepEvalBackend (deepmd/pretrained/deep_eval.py#L55-L60) is an eighth DeepEvalBackend subclass not named in the PR body's "seven adapters"; it is a pure delegator so it inherits the fix, but it is worth a mention so nobody assumes it was audited and missed.

Unrelated to the diff: the Read the Docs check is failing on this head. I could not read the build log (the page is not publicly accessible), and #5908 fails it too while other open PRs pass, so it may well be infrastructure rather than yours -- flagging only so it is not overlooked.

Reviewed head 1709902d73dc6dfb92795b0cb956d3716d9c5df0 with three independent full subagent reviews. No actionable code findings were identified. The shared normalization preserves the documented shorthand semantics and historical backend ABI across the reviewed adapters. The PR is currently conflicting, and the Read the Docs and pre-commit checks should be rerun after rebasing. Because this PR was opened by the active njzjz-bot account, this is a comment-only review rather than a self-approval event.

The quota is about to reset, so I am concentrating the remaining token budget on these reviews.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh